'''
  This class represent an AVL Tree, an AVL tree is a self-balancing
  binary search tree, and it was the first such data structure to 
  be invented. In an AVL tree, the heights of the two child subtrees 
  of any node differ by at most one; if at any time they differ 
  by more than one, re-balancing is done to restore this property.

  Author: Samuel T. C. Santos
  Date  : 2013/04/19

'''

from bst import BST
from node import Node

class AVL(BST):

	def __init__(self):
		self.root = Node(None)
		self.root.setParent(Node(None))

	def insert(self, element):
		if (self.root.isEmpty()):
			self.root.setData(element)
			self.root.setLeft(Node(None))
			self.root.setRight(Node(None))
			self.root.setParent(Node(None))
			self.root.getLeft().setParent(self.root)
			self.root.getRight().setParent(self.root)
		else:
			root = self.getRoot()
			self.insertNode(element,root)		

	def insertNode(self, element, node):
		if element < node.getData():
			if node.getLeft().isEmpty():
				node.getLeft().setParent(node)
				node.getLeft().setData(element)
				node.getLeft().setLeft(Node(None))
				node.getLeft().setRight(Node(None))
			else:
				self.insertNode(element, node.getLeft())
		else:
			if node.getRight().isEmpty():
				node.getRight().setParent(node)
				node.getRight().setData(element)
				node.getRight().setLeft(Node(None))
				node.getRight().setRight(Node(None))
			else:
				self.insertNode(element, node.getRight())

		self.rebalance(node)


	def remove(self, element):
		node = self.search(element)

		if not node.isEmpty():
			self.removeNode(node)

	def removeNode(self, node):
		if node.getLeft().isEmpty() and node.getRight().isEmpty():
			node.setData(None)
			node.setLeft(Node(None))
			node.setRight(Node(None))
			self.rebalanceUp(node)

		elif node.hasOneChild():
			if not node.equals(self.getRoot()):
			
				if node.isLeftChild():
					if not node.getLeft().isEmpty():
						node.getParent().setLeft(node.getLeft())
					else:
						node.getParent().setLeft(node.getRight())
				else:
					if not node.getLeft().isEmpty():
						node.getParent().setRight(node.getLeft())
					else:
						node.getParent().setRight(node.getRight())
			else:
				if not node.getRight().isEmpty():
					self.root = node.getRight()
				else:
					self.root = node.getLeft()

				self.root.setParent(Node(None))

			self.rebalanceUp(node)
		else:
			successor = self.successorNode(node)
			node.setData(successor.getData())
			self.removeNode(successor)

	'''
	 This method do bottom-up rebalance tree. 
	'''
	def rebalanceUp(self, node):
		parent = node.getParent()

		while not parent.isEmpty():
			self.rebalance(parent)
			parent = parent.getParent()


	'''
 	 This method re-balance the tree.
	'''
	def rebalance(self, node):
		
		if self.balanceFactorNode(node) > 1:
			print("Unbalanced")
			if self.balanceFactorNode(node.getLeft()) > 0:
				self.rightRotationNode(node)
			else:
				self.leftRotationNode(node.getLeft())
				self.rightRotationNode(node)
		elif self.balanceFactorNode(node) < -1:
			print("Unbalanced")
			if self.balanceFactorNode(node.getRight()) < 0:
				self.leftRotationNode(node)
			else:
				self.rightRotationNode(node.getRight())
				self.leftRotationNode(node)
		else:
			print ("balanced")

	'''
    The balance factor return the difference between 
    the left height and right height.

	'''
	def balanceFactor(self, element):
		node = self.search(element)

		if not node.isEmpty():
			return self.balanceFactorNode(node)
		else:
			return 0 

	def balanceFactorNode(self, node):
		return self.heightNode(node.getLeft()) - self.heightNode(node.getRight())

	'''
	 
	 The rotation operation on a AVL.

	'''

	def leftRotation(self, element):
		node = self.search(element)

		if not node.isEmpty():
			self.leftRotationNode(node)


	def leftRotationNode(self, node):
		
		pivot = node.getRight()
		node.setRight(pivot.getLeft())
		
		if not pivot.isEmpty():
			pivot.getLeft().setParent(node)
		
		pivot.getParent().setParent(node.getParent())
		
		if node.getParent().isEmpty():
			pivot.setParent(Node(None))
			self.root = pivot
		elif node.equals(node.getParent().getLeft()):
			node.getParent().setLeft(pivot)
		else:
			node.getParent().setRight(pivot)
		
		pivot.setLeft(node)
		node.setParent(pivot)

	'''
		This method do an right rotate.
		
	'''
	def rightRotation(self, element):
		node = self.search(element)

		if not node.isEmpty():
			self.rightRotationNode(node)

	def rightRotationNode(self, node):

		pivot = node.getLeft()
		node.setLeft(pivot.getRight())

		if not pivot.isEmpty():
			pivot.getRight().setParent(node)

		pivot.getParent().setParent(node.getParent())

		if node.getParent().isEmpty():
			pivot.setParent(Node(Node))
			self.root = pivot
		elif node.equals(node.getParent().getRight()):
			node.getParent().setRight(pivot)
		else:
			node.getParent().setLeft(pivot)

		pivot.setRight(node)
		node.setParent(pivot)

